import json
import time
import math
from datetime import datetime
from pynput import keyboard, mouse
from pynput.keyboard import Key, KeyCode
from collections import deque
import threading


class InputRecorder:
    def __init__(self):
        self.recording = False
        self.start_time = 0
        self.events = []
        self.session_id = datetime.now().strftime("%Y-%m-%d-%H%M%S")
        self.video_name = f"recording_{self.session_id}"

        # 视角追踪相关变量
        self.current_view_direction = "Neutral"
        self.view_change_time = 0
        self.mouse_moves = []

        # 运动学模拟相关变量
        self.kinematic_data = []
        self.last_kinematic_update = 0
        self.current_position = {"x": 0.0, "y": 0.0, "z": 0.0}
        self.current_velocity = {"x": 0.0, "y": 0.0, "z": 0.0}
        self.current_acceleration = {"x": 0.0, "y": 0.0, "z": 0.0}
        self.active_movement_keys = set()

        # 状态追踪相关变量
        self.current_states = set()
        self.state_history = []

        # 初始化监听器
        self.keyboard_listener = keyboard.Listener(on_press=self.on_key_press, on_release=self.on_key_release)
        self.mouse_listener = mouse.Listener(on_click=self.on_click, on_scroll=self.on_scroll, on_move=self.on_move)

    def on_move(self, x, y):
        """处理鼠标移动事件 - 用于视角追踪"""
        if self.recording:
            current_time = time.time() - self.start_time
            self.mouse_moves.append({
                "x": x,
                "y": y,
                "timestamp": round(current_time, 3)
            })

            # 限制存储的鼠标移动数据量
            if len(self.mouse_moves) > 1000:
                self.mouse_moves = self.mouse_moves[-500:]

    def on_key_press(self, key):
        """处理按键按下事件"""
        try:
            # 检查是否是F8键
            if hasattr(key, 'name') and key.name == 'f8':
                self.toggle_recording()
                return

            if self.recording:
                current_time = time.time() - self.start_time
                key_name = key.char if hasattr(key, 'char') and key.char else str(key).replace('Key.', '')

                # 添加到活动按键集合
                self.active_movement_keys.add(key_name.lower())

                # 检查状态变化（跳跃、蹲下等）
                self.check_state_changes(key_name, True, current_time)

                event = {
                    "type": "key_press",
                    "key": key_name,
                    "timestamp": round(current_time, 3)
                }
                self.events.append(event)
                print(f"记录按键按下: {key_name} at {current_time:.3f}s")

        except Exception as e:
            print(f"处理按键按下时出错: {e}")

    def on_key_release(self, key):
        """处理按键释放事件"""
        if self.recording:
            try:
                # 忽略F8键的释放事件
                if hasattr(key, 'name') and key.name == 'f8':
                    return

                current_time = time.time() - self.start_time
                key_name = key.char if hasattr(key, 'char') and key.char else str(key).replace('Key.', '')

                # 从活动按键集合中移除
                if key_name.lower() in self.active_movement_keys:
                    self.active_movement_keys.remove(key_name.lower())

                # 检查状态变化
                self.check_state_changes(key_name, False, current_time)

                event = {
                    "type": "key_release",
                    "key": key_name,
                    "timestamp": round(current_time, 3)
                }
                self.events.append(event)

                # 更新运动学数据
                self.update_kinematic_data(current_time)

            except Exception as e:
                print(f"处理按键释放时出错: {e}")

    def check_state_changes(self, key, is_press, timestamp):
        """检查并记录状态变化（跳跃、蹲下等）"""
        state_map = {
            'space': 'JUMPING',
            'ctrl': 'CROUCHING',
            'shift': 'SPRINTING',
            'alt': 'WALKING'
        }

        key_lower = key.lower()
        if key_lower in state_map:
            state = state_map[key_lower]
            if is_press and state not in self.current_states:
                self.current_states.add(state)
                self.state_history.append({
                    "timestamp": round(timestamp, 3),
                    "state": state,
                    "action": "start"
                })
            elif not is_press and state in self.current_states:
                self.current_states.remove(state)
                self.state_history.append({
                    "timestamp": round(timestamp, 3),
                    "state": state,
                    "action": "end"
                })

    def on_click(self, x, y, button, pressed):
        """处理鼠标点击事件"""
        if self.recording:
            current_time = time.time() - self.start_time
            action = "按下" if pressed else "释放"
            event = {
                "type": f"mouse_{'press' if pressed else 'release'}",
                "button": str(button).replace('Button.', ''),
                "x": x,
                "y": y,
                "timestamp": round(current_time, 3)
            }
            self.events.append(event)

            # 更新运动学数据（射击后坐力模拟）
            if pressed and str(button).replace('Button.', '') == 'left':
                self.current_velocity['x'] += (x - 960) / 1000  # 简单模拟后坐力
                self.current_velocity['y'] += (y - 540) / 1000

            print(f"记录鼠标{action}: {button} at {current_time:.3f}s")

    def on_scroll(self, x, y, dx, dy):
        """处理鼠标滚动事件"""
        if self.recording:
            current_time = time.time() - self.start_time

            # 确定视角方向
            new_direction = "U" if dy > 0 else "D"
            if new_direction != self.current_view_direction:
                # 记录前一个方向段的结束
                if self.view_change_time > 0:
                    self.events.append({
                        "type": "view_direction",
                        "direction": self.current_view_direction,
                        "start_time": round(self.view_change_time, 3),
                        "end_time": round(current_time, 3)
                    })

                # 更新当前方向
                self.current_view_direction = new_direction
                self.view_change_time = current_time

            event = {
                "type": "mouse_scroll",
                "dx": dx,
                "dy": dy,
                "x": x,
                "y": y,
                "timestamp": round(current_time, 3)
            }
            self.events.append(event)
            print(f"记录鼠标滚动: dx={dx}, dy={dy} at {current_time:.3f}s")

    def update_kinematic_data(self, current_time):
        """更新运动学数据（基于按键状态模拟）"""
        if current_time - self.last_kinematic_update < 0.1:  # 每100ms更新一次
            return

        # 基于活动按键更新速度和加速度
        base_speed = 2.5
        acceleration_rate = 1.2

        # 重置加速度
        self.current_acceleration = {"x": 0.0, "y": 0.0, "z": 0.0}

        # 根据按键更新加速度
        if 'w' in self.active_movement_keys:
            self.current_acceleration['z'] = acceleration_rate
        elif 's' in self.active_movement_keys:
            self.current_acceleration['z'] = -acceleration_rate

        if 'a' in self.active_movement_keys:
            self.current_acceleration['x'] = -acceleration_rate
        elif 'd' in self.active_movement_keys:
            self.current_acceleration['x'] = acceleration_rate

        # 更新速度
        for axis in ['x', 'y', 'z']:
            self.current_velocity[axis] += self.current_acceleration[axis] * 0.1
            # 速度衰减（模拟阻力）
            self.current_velocity[axis] *= 0.9

        # 更新位置
        for axis in ['x', 'y', 'z']:
            self.current_position[axis] += self.current_velocity[axis] * 0.1

        # 记录运动学数据
        self.kinematic_data.append({
            "timestamp": round(current_time, 3),
            "position": {k: round(v, 2) for k, v in self.current_position.items()},
            "velocity": {k: round(v, 1) for k, v in self.current_velocity.items()},
            "acceleration": {k: round(v, 1) for k, v in self.current_acceleration.items()}
        })

        self.last_kinematic_update = current_time

    def toggle_recording(self):
        """切换录制状态"""
        if not self.recording:
            # 开始录制
            self.recording = True
            self.start_time = time.time()
            self.events = []
            self.mouse_moves = []
            self.kinematic_data = []
            self.state_history = []
            self.current_view_direction = "Neutral"
            self.view_change_time = 0
            self.active_movement_keys = set()
            self.current_states = set()

            # 初始化运动学数据
            self.current_position = {"x": 0.0, "y": 0.0, "z": 0.0}
            self.current_velocity = {"x": 0.0, "y": 0.0, "z": 0.0}
            self.current_acceleration = {"x": 0.0, "y": 0.0, "z": 0.0}
            self.last_kinematic_update = 0

            print("开始记录输入...")
            print("按下F8键停止记录")

            # 启动运动学数据更新线程
            self.kinematic_update_thread = threading.Thread(target=self.kinematic_update_loop)
            self.kinematic_update_thread.daemon = True
            self.kinematic_update_thread.start()

        else:
            # 停止录制
            self.recording = False
            duration = time.time() - self.start_time

            # 确保记录最后的视角段
            if self.view_change_time > 0:
                self.events.append({
                    "type": "view_direction",
                    "direction": self.current_view_direction,
                    "start_time": round(self.view_change_time, 3),
                    "end_time": round(duration, 3)
                })

            print(f"停止记录输入，时长: {duration:.2f}秒")
            self.export_to_json(duration)

    def kinematic_update_loop(self):
        """运动学数据更新循环（在独立线程中运行）"""
        while self.recording:
            current_time = time.time() - self.start_time
            self.update_kinematic_data(current_time)
            time.sleep(0.1)  # 每100ms更新一次

    def export_to_json(self, duration):
        """将记录的数据导出为JSON格式"""
        # 处理各种事件数据
        movement_events = self.process_movement_events()
        view_events = self.process_view_events()
        kinematic_data = self.kinematic_data
        additional_states = self.process_additional_states()

        # 构建JSON数据结构
        data = {
            "game_metadata": {
                "session_id": self.session_id,
                "video_name": self.video_name,
                "start_time": datetime.fromtimestamp(self.start_time).isoformat(),
                "duration_seconds": round(duration, 3)
            },
            "movement_events": movement_events,
            "view_events": view_events,
            "kinematic_data": kinematic_data,
            "additional_states": additional_states
        }

        # 生成文件名
        filename = f"input_recording_{self.session_id}.json"

        # 写入JSON文件
        with open(filename, 'w') as f:
            json.dump(data, f, indent=2)

        print(f"记录已导出到: {filename}")
        print(f"移动事件: {len(movement_events)} 条")
        print(f"视角事件: {len(view_events)} 条")
        print(f"运动学数据: {len(kinematic_data)} 条")
        print(f"附加状态: {len(additional_states)} 条")

    def process_movement_events(self):
        """处理移动事件"""
        movement_events = []
        key_states = {}

        for event in self.events:
            if event['type'] in ['key_press', 'key_release']:
                key = event['key']
                is_press = event['type'] == 'key_press'

                # 只处理移动相关按键
                movement_keys = ['w', 'a', 's', 'd', 'up', 'down', 'left', 'right']
                if key.lower() in movement_keys:
                    if is_press:
                        key_states[key] = event['timestamp']
                    else:
                        if key in key_states:
                            movement_events.append({
                                "direction": key.upper(),
                                "start_time": round(key_states[key], 3),
                                "end_time": round(event['timestamp'], 3)
                            })
                            del key_states[key]

        return movement_events

    def process_view_events(self):
        """处理视角事件"""
        view_events = []

        # 从事件中提取视角方向变化
        for event in self.events:
            if event['type'] == 'view_direction':
                view_events.append({
                    "direction": event['direction'],
                    "start_time": event['start_time'],
                    "end_time": event['end_time']
                })

        # 如果没有视角事件，添加一个默认的中性视角
        if not view_events:
            view_events.append({
                "direction": "Neutral",
                "start_time": 0.0,
                "end_time": round(time.time() - self.start_time, 3)
            })

        return view_events

    def process_additional_states(self):
        """处理附加状态"""
        additional_states = []
        current_states = {}

        for event in self.state_history:
            state = event['state']
            timestamp = event['timestamp']
            action = event['action']

            if action == 'start':
                current_states[state] = timestamp
            elif action == 'end' and state in current_states:
                additional_states.append({
                    "timestamp": round(timestamp, 3),
                    "state": state
                })
                del current_states[state]

        # 添加仍在进行中的状态
        for state, start_time in current_states.items():
            additional_states.append({
                "timestamp": round(start_time, 3),
                "state": state
            })

        # 按时间戳排序
        additional_states.sort(key=lambda x: x['timestamp'])

        return additional_states

    def start(self):
        """启动监听器"""
        print("输入记录器已启动...")
        print("按下F8键开始记录")
        self.keyboard_listener.start()
        self.mouse_listener.start()
        self.keyboard_listener.join()
        self.mouse_listener.join()


if __name__ == "__main__":
    recorder = InputRecorder()
    recorder.start()
